from flask import Flask, Response
import cv2
from ultralytics import YOLO
import threading

app = Flask(__name__)


class RISyolo:
    def __init__(self, model_path):
        self.camera = None
        self.model = YOLO(model_path, verbose=False)
        self.latest_frame = None
        self.frame_lock = threading.Lock()

    def set_camera(self, camera_index=4):
        self.camera = cv2.VideoCapture(camera_index)

    def release_camera(self):
        self.camera.release()

    def capture_frame(self):
        while True:
            ret, frame = self.camera.read()
            if not ret:
                break
            results = self.model(frame,verbose=False)
            annotated_frame = results[0].plot()
            with self.frame_lock:
                self.latest_frame = frame
        self.release_camera()

    def get_original_frame(self):
        with self.frame_lock:
            return self.latest_frame

    def get_encoded_frame(self):
        with self.frame_lock:
            ret, buffer = cv2.imencode('.jpg', self.latest_frame)
            frame = buffer.tobytes()
            return frame

    def start(self):
        threading.Thread(target=self.capture_frame).start()


camera = RISyolo("./models/yolo11n.pt")
camera.set_camera()


def generate_frames():
    while True:
        frame = camera.get_encoded_frame()
        print(frame)
        if frame is None:
            pass
        else:
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')


@app.route('/')
def index():
    return "<h1>Video Stream</h1><img src='/video_feed'>"


@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')


camera.start()